import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.autograd import Function

class STEMask(Function):
    @staticmethod
    def forward(ctx, w, m):
        return w * m

    @staticmethod
    def backward(ctx, grad_out):
        return grad_out, None

class STEConv2d(nn.Conv2d):
    # magnitude-based pruning
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.register_buffer('mask', torch.ones_like(self.weight))
    
    def forward(self, x):
        w = STEMask.apply(self.weight, self.mask)
        return F.conv2d(x, w, self.bias, self.stride, self.padding, self.dilation, self.groups)
    
    @classmethod
    def from_dense(cls, mod: nn.Conv2d):
        new_mod = cls(
            in_channels = mod.in_channels,
            out_channels = mod.out_channels,
            kernel_size = mod.kernel_size,
            bias = mod.bias is not None,
            stride = mod.stride,
            padding = mod.padding,
            dilation = mod.dilation,
            groups = mod.groups
        )
        with torch.no_grad():
            new_mod.weight.copy_(mod.weight)
            if mod.bias is not None:
                new_mod.bias.copy_(mod.bias)
        return new_mod.to(mod.weight.device)

class STELinear(nn.Linear):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.register_buffer('mask', torch.empty_like(self.weight))
    
    def forward(self, x):
        w = STEMask.apply(self.weight, self.mask)
        return F.linear(x, w, self.bias)
    
    @classmethod
    def from_dense(cls, mod: nn.Linear):
        new_mod = cls(
            in_features = mod.in_features,
            out_features = mod.out_features,
            bias = mod.bias is not None,
        )
        with torch.no_grad():
            new_mod.weight.copy_(mod.weight)
            if mod.bias is not None:
                new_mod.bias.copy_(mod.bias)
        return new_mod.to(mod.weight.device)

mapping = {
    nn.Conv2d: STEConv2d,
    nn.Linear: STELinear
}